import networkx as nx
import numpy as np
import random
from iteration_helpers import opinion_update, dissim_deg, delete_tie, add_tie, update


################################ Power Broker ################################
def add_power_broker(G, broker):
  total_degree = sum(dict(G.degree()).values())
  num_nodes = G.number_of_nodes()
  k = int(total_degree / num_nodes)

  opinions = [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]
  identity_values = [-1, 1]
  identity = np.random.choice(identity_values) 
  beta = np.clip(np.random.normal(0.5, 0.15), 0, 1)
  alpha = np.clip(np.random.normal(1, 0.3), 0, 2)

  G.add_node(broker, opinions = opinions, identity = identity, 
             beta = beta, alpha = alpha)

  for node in range(k):
      non_neighbors = list(nx.non_neighbors(G, broker))
      if not non_neighbors:
         return
      target = random.choice(non_neighbors)
      G.add_edge(broker, target)

def power_broker_update(G, broker_name):
  if random.random() >= 0.5:
    # Normal update to the graph
    random_node = random.choice(list(G.nodes))
    opinion_update(G, random_node)
    delete_tie(G, random_node)
    add_tie(G, random_node)
  
  else:
    # Parallel broker update
    delete_tie(G, broker_name)
    add_tie(G, broker_name)


################################ Structural Broker ################################
def add_structural_broker(G, broker):
  initial_betweenness = nx.betweenness_centrality(G)
  sorted_nodes = sorted(initial_betweenness, key=initial_betweenness.get, reverse=True)
  total_degree = sum(dict(G.degree()).values())
  num_nodes = G.number_of_nodes()
  k = int(total_degree / num_nodes)
  top_k_nodes = sorted_nodes[:k]

  opinions = np.random.normal(0, 0.15, 10)
  opinions = [max(-1, min(1, x)) for x in opinions]
  identity_values = [-1, 1]
  identity = np.random.choice(identity_values) 
  beta = np.clip(np.random.normal(0.5, 0.15), 0, 1)
  alpha = np.clip(np.random.normal(1, 0.3), 0, 2)

  G.add_node(broker, opinions = opinions, identity = identity, 
             beta = beta, alpha = alpha)

  for node in top_k_nodes:
      G.add_edge(broker, node)

def structural_broker_update(G, broker_name):
  if random.random() >= 0.5:
    # Normal update to the graph
    random_node = random.choice(list(G.nodes))
    opinion_update(G, random_node)
    delete_tie(G, random_node)
    add_tie(G, random_node)
  
  else:
    # Parallel broker update
    betweenness = nx.betweenness_centrality(G)
    sorted_nodes = sorted(betweenness, key = betweenness.get, reverse=True)
    opinion_update(G, broker_name) # Update broker's opinions based on its neighbors
    delete_tie(G, broker_name)
    # add_tie
    edge_added = False
    for node in sorted_nodes:
      if edge_added:
        break
      else:
        if G.has_edge(node, broker_name):
          continue
        dis = dissim_deg(G, broker_name, node)
        broker_alpha = G.nodes[broker_name]['alpha']
        if dis < broker_alpha:
          G.add_edge(broker_name, node)
          edge_added = True
        elif dis == broker_alpha:
          if random.random() < 0.5:
            G.add_edge(broker_name, node)
            edge_added = True

################################ Social Broker ################################
def add_social_broker(G, broker):
  total_degree = sum(dict(G.degree()).values())
  num_nodes = G.number_of_nodes()
  k = int(total_degree / num_nodes)

  opinions = np.random.normal(0, 0.15, 10)
  opinions = [max(-1, min(1, x)) for x in opinions]
  identity_values = [-1, 1]
  identity = np.random.choice(identity_values) 
  beta = np.clip(np.random.normal(0.5, 0.15), 0, 1)
  alpha = np.clip(np.random.normal(1, 0.3), 0, 2)


  G.add_node(broker, opinions = opinions, identity = identity, 
             beta = beta, alpha = alpha)

  for node in range(k):
      non_neighbors = list(nx.non_neighbors(G, broker))
      if not non_neighbors:
         return
      target = random.choice(non_neighbors)
      G.add_edge(broker, target)

def hetero_opinion_update(G, node):
  # Fetch neighbors, identity, and opinions of the node
  neighbors = list(G[node])
  identity = G.nodes[node]['identity']
  opinions = G.nodes[node]['opinions']
  beta = G.nodes[node]['beta']

  # Initialize variables for co-partisan and opp-partisan opinions
  opp_opinions = np.zeros_like(opinions)
  opp_count = 0
  co_opinions = np.zeros_like(opinions)
  co_count = 0

  # Aggregate opinions from neighbors based on their identities
  for neighbor in neighbors:
    neighbor_identity = G.nodes[neighbor]['identity']
    neighbor_opinion = G.nodes[neighbor]['opinions']
    if neighbor_identity == identity:
      co_count += 1
      co_opinions += neighbor_opinion
    else: 
      opp_count += 1
      opp_opinions += neighbor_opinion

  # Compute average opinions, handling cases where counts are zero
  if co_count > 0:
    co_opinions /= co_count
  if opp_count > 0:
    opp_opinions /= opp_count
  
  # Update the node's opinions based on beta and neighbors' opinions (heterophily)
  new_opinions = opinions - beta * co_opinions + beta * opp_opinions
  new_opinions = np.tanh(new_opinions)
  G.nodes[node]['opinions'] = new_opinions

def hetero_delete_tie(G, node):
  # Fetch neighbors
  neighbors = list(G[node])
  if not neighbors:
        return  # No neighbors to delete

  # Randomly choose one of the node's neighbors and calculate dissimilarity degree
  target = random.choice(neighbors)
  dis = dissim_deg(G, node, target)
  alpha = G.nodes[node]['alpha']

  # Delete the tie if dissimilarity degree < alpha
  if dis < alpha: 
    G.remove_edge(node, target)
  
  # Delete with 50% chance if dissimilarity degree == alpha
  if dis == alpha:
    if random.random() < 0.5:
      G.remove_edge(node, target)

def hetero_add_tie(G, node):
  # Randomly choose a non-neighbor and calculate dissimilarity degree
  non_neighbors = list(nx.non_neighbors(G, node))
  if not non_neighbors:
        return  # No potential nodes to connect
  target = random.choice(non_neighbors)
  dis = dissim_deg(G, node, target)
  alpha = G.nodes[node]['alpha']

  # Add the tie if dissimilarity degree < alpha
  if dis > alpha:
    G.add_edge(node, target)

  # Add with 50% cance if dissimilarity degree == alpha
  if dis == alpha:
    if random.random() < 0.5:
      G.add_edge(node, target)

def social_broker_update(G, broker_name):
  if random.random() >= 0.5:
    # Normal update to the graph
    random_node = random.choice(list(G.nodes))
    opinion_update(G, random_node)
    delete_tie(G, random_node)
    add_tie(G, random_node)

  else:
    # heterophily-based update rules
    hetero_opinion_update(G, broker_name) # Update broker's opinions based on its neighbors
    hetero_delete_tie(G, broker_name)
    hetero_add_tie(G, broker_name)






################################ Optimal Broker ################################
def add_optimal_broker(G, broker):
  initial_betweenness = nx.betweenness_centrality(G)
  sorted_nodes = sorted(initial_betweenness, key=initial_betweenness.get, reverse=True)
  total_degree = sum(dict(G.degree()).values())
  num_nodes = G.number_of_nodes()
  k = int(total_degree / num_nodes)
  top_k_nodes = sorted_nodes[:k]

  opinions = [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]
  identity_values = [-1, 1]
  identity = np.random.choice(identity_values) 
  beta = np.clip(np.random.normal(0.5, 0.15), 0, 1)
  alpha = np.clip(np.random.normal(1, 0.3), 0, 2)

  G.add_node(broker, opinions = opinions, identity = identity, 
             beta = beta, alpha = alpha)

  for node in top_k_nodes:
      G.add_edge(broker, node)

def optimal_opinion_update(G, node):
  # Fetch neighbors, identity, and opinions of the node
  neighbors = list(G[node])
  identity = G.nodes[node]['identity']
  opinions = G.nodes[node]['opinions']
  beta = G.nodes[node]['beta']

  # Initialize variables for co-partisan and opp-partisan opinions
  opp_opinions = np.zeros_like(opinions)
  opp_count = 0
  co_opinions = np.zeros_like(opinions)
  co_count = 0

  # Aggregate opinions from neighbors based on their identities
  for neighbor in neighbors:
    neighbor_identity = G.nodes[neighbor]['identity']
    neighbor_opinion = G.nodes[neighbor]['opinions']
    if neighbor_identity == identity:
      co_count += 1
      co_opinions += neighbor_opinion
    else: 
      opp_count += 1
      opp_opinions += neighbor_opinion

  # Compute average opinions, handling cases where counts are zero
  if co_count > 0:
    co_opinions /= co_count
  if opp_count > 0:
    opp_opinions /= opp_count
  
  # Update the node's opinions based on beta and neighbors' opinions (heterophily)
  new_opinions = opinions - beta * co_opinions + beta * opp_opinions
  new_opinions = np.tanh(new_opinions)
  G.nodes[node]['opinions'] = new_opinions

def optimal_delete_tie(G, node):
  # Fetch neighbors
  neighbors = list(G[node])
  if not neighbors:
        return  # No neighbors to delete

  # Randomly choose one of the node's neighbors and calculate dissimilarity degree
  target = random.choice(neighbors)
  dis = dissim_deg(G, node, target)
  alpha = G.nodes[node]['alpha']

  # Delete the tie if dissimilarity degree < alpha
  if dis < alpha: 
    G.remove_edge(node, target)
  
  # Delete with 50% chance if dissimilarity degree == alpha
  if dis == alpha:
    if random.random() < 0.5:
      G.remove_edge(node, target)

def optimal_add_tie(G, broker_name):
  # Extract the node with the highest betweenness centrality
  betweenness = nx.betweenness_centrality(G)
  sorted_nodes = sorted(betweenness, key = betweenness.get, reverse=True)
  
  edge_added = False
  for node in sorted_nodes:
    if not edge_added:
      if G.has_edge(node, broker_name):
         continue
      dis = dissim_deg(G, broker_name, node)
      alpha = G.nodes[broker_name]['alpha']
      # add tie based on heterophily
      if dis > alpha:
        G.add_edge(broker_name, node)
        edge_added = True
      elif dis == alpha:
         if random.random() < 0.5:
            G.add_edge(broker_name, node)
            edge_added = True
    if edge_added:
       break
    
def optimal_broker_update(G, broker_name):
  if random.random() >= 0.5: 
    random_node = random.choice(list(G.nodes))
    opinion_update(G, random_node)
    delete_tie(G, random_node)
    add_tie(G, random_node)
  
  else:
    optimal_opinion_update(G, broker_name) # Update broker's opinions based on its neighbors
    optimal_delete_tie(G, broker_name)
    optimal_add_tie(G, broker_name)
